      
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from sklearn.metrics import confusion_matrix, classification_report
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from torchinfo import summary
import cv2
import os


def create_layer(input_dim, output_dim, LeakyReLU, dropout_rate):
    return nn.Linear(input_dim, output_dim), nn.BatchNorm1d(output_dim), nn.LeakyReLU(LeakyReLU), nn.Dropout(dropout_rate)

class OptimizedMLP(nn.Module):
    def __init__(self, Flatten,layers):
        super(OptimizedMLP, self).__init__()

        self.flatten = Flatten
        self.layers = nn.ModuleList(layers)

    def forward(self, x):
        x = self.flatten(x)
        for layer in self.layers:
            x = layer(x)
        
        return x


# 权重初始化函数
def weights_init(m): #用于初始化模型的权重。
    if isinstance(m, nn.Linear): #检查是否是全连接层。
        nn.init.kaiming_uniform_(m.weight) #使用 Kaiming 均匀分布初始化权重，这是一种常用的初始化方法，可以加速模型的收敛。
        if m.bias is not None:     #检查是否有偏置项。
            nn.init.zeros_(m.bias) #将偏置项初始化为零。

# 模型评估函数，计算模型在测试集上的准确率
def evaluate_model(model, test_loader, device): 
    model.eval() # 将模型设置为评估模式。
    correct = 0     # 初始化正确预测的样本数。
    total = 0 # 初始化总样本数量。

    with torch.no_grad():   # 在评估模式下，不需要计算梯度，可以节省内存和计算资源。
        for images, labels in test_loader:  # 遍历测试集的每个批次。
            images = images.view(-1, 28*28).to(device)  # 展平输入图像
            labels = labels.to(device)  # 将标签移动到设备上
            outputs = model(images)            # 前向传播，得到模型的输出
            _, predicted = torch.max(outputs.data, 1)           # 获取预测结果
            total += labels.size(0)         # 累加当前批次的样本数量
            correct += (predicted == labels).sum().item()  # 累加当前批次的正确预测数量

    accuracy = 100 * correct / total    # 计算准确率
    return accuracy         # 返回准确率。

# 加载模型函数
def load_model(model_path, device):
    device = torch.device(device) #检查是否有可用的 GPU，如果有，则使用 GPU。
    model = torch.load(model_path, map_location=device)
    model.eval()
    return model

# 识别数字函数
def recognize_digits(model, processed_digits, device):
    results = []                    # 初始化结果列表。
    with torch.no_grad():                   
        for digit in processed_digits:  # 遍历每个数字。
            tensor_digit = torch.from_numpy(digit).unsqueeze(0).float().to(device)
            tensor_digit = tensor_digit.view(-1, 28*28)  # 展平输入，将输入张量展平为一维向量，形状为(1, 784)，以适应模型的输入要求。
            # 标准化处理，与训练时的归一化一致
            tensor_digit = (tensor_digit - 0.5) / 0.5 #对输入张量进行标准化处理，使其像素值范围在[-1, 1]之间，与训练时的归一化方式保持一致。
            output = model(tensor_digit) #将标准化后的输入张量输入到模型中进行前向传播，得到模型的输出。
            _, predicted = torch.max(output.data, 1) #获取模型输出中概率最大的类别索引
            results.append(predicted.item())
    return results
    
# def test_q1(path,opt):
#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#     model = opt.to(device)
#     model.eval()
    
#     # 在这里，我们将使用测试集中的前1个样本作为示例
#     test_dataset = datasets.MNIST(root=path, train=False, download=True, transform=transforms.ToTensor())
#     tensor_digit, label = test_dataset[0]       # 第一个样本
#     print('样本图像维度为: ', tensor_digit.shape)
#     print('模型的结果信息: ')
#     # output = model(tensor_digit) 
#     # print('output:', output)
#     summary(model, tensor_digit.shape, col_names=["input_size", "output_size"])

def test_q1(opt):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = opt.to(device)
    model.eval()

    # 创建一个随机输入张量
    tensor_digit = torch.randn(1, 28, 28).to(device)

    print('图像维度为: ', tensor_digit.shape)
    # output = model(tensor_digit) 
    # print('output:', output)
    print('模型的结构信息: ')
    # summary(model, tensor_digit.shape)
    summary(model, tensor_digit.shape, col_names=["input_size", "output_size"])

# 预处理图像
def preprocess_image(image_path, save_intermediate=False, output_dir=None, image_name=None):
    image_name = os.path.splitext(os.path.basename(image_name))[0]
    if save_intermediate:
        # 创建保存中间结果的目录
        image_output_dir = os.path.join(output_dir, image_name)
        os.makedirs(image_output_dir, exist_ok=True)
    else:
        image_output_dir = None
    # 读取图像
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"无法打开图像文件：{image_path}")

    # 转换为灰度图像
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) #将 BGR 格式的图像转换为灰度图像。
    if save_intermediate:
        os.makedirs(image_output_dir, exist_ok=True)
        # 保存灰度图像，灰度化的目的：减少数据维度（从 3 通道变为 1 通道），同时保留图像的主要特征，便于后续处理。
        cv2.imwrite(os.path.join(image_output_dir, f"{image_name}_0_input.png"), img)

    # 使用自适应阈值进行二值化，将灰度图像转换为二值图像。
    #二值化的目的：将图像中的数字与背景分离，便于后续的轮廓检测和数字分割。
    thresh = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 15, 10
    )
    if save_intermediate:
        # 保存二值化后的图像
        cv2.imwrite(os.path.join(image_output_dir, f"{image_name}_1_binarization.png"), thresh)
    return thresh


# 分割数字
#从预处理后的二值图像中分割出单个数字的图像。
def segment_digits(processed_image, save_intermediate=False, output_dir=None, image_name=None):
    image_name = os.path.splitext(os.path.basename(image_name))[0]
    if save_intermediate:
        # 创建保存中间结果的目录
        image_output_dir = os.path.join(output_dir, image_name)
        os.makedirs(image_output_dir, exist_ok=True)
    else:
        image_output_dir = None
    # 寻找图像中的轮廓
    contours, _ = cv2.findContours(
        processed_image.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )

    digit_imgs = []  # 存储每个数字的图像
    bounding_boxes = []  # 存储每个数字的边界框
    image_height, image_width = processed_image.shape #获取输入图像的高度和宽度

    if save_intermediate:
        # 如果需要保存中间结果，转换二值图像为彩色图像，以绘制彩色边界框
        image_with_boxes = cv2.cvtColor(processed_image.copy(), cv2.COLOR_GRAY2BGR)

    for cnt in contours:#：对每个轮廓计算边界框，并过滤掉过小的区域（认为是噪声）。
        # 获取轮廓的边界框，返回边界框的左上角坐标 (x, y) 和宽高 (w, h)。
        x, y, w, h = cv2.boundingRect(cnt)

        # 过滤掉太小的区域，认为是噪声
        if w > image_width / 20 or h > image_height / 20:
            #过滤条件：如果边界框的宽度或高度大于图像宽度或高度的 1/20，则认为该区域可能是数字。
            bounding_boxes.append((x, y, w, h))
            # if save_intermediate:
            #     # 在图像上绘制矩形框，用于标记检测到的数字区域。
            #     cv2.rectangle(
            #         image_with_boxes, (x, y), (x + w, y + h), (0, 255, 0), 2
            #     )

    if not bounding_boxes:
        #如果没有检测到符合条件的轮廓，返回整个输入图像作为分割结果。
        print("Warning: No contours found after filtering. Returning the entire image.")
        return [processed_image]

    # 按照从左到右的顺序排序边界框
    bounding_boxes = sorted(bounding_boxes, key=lambda box: box[0])

    for idx, box in enumerate(bounding_boxes): #根据边界框提取每个数字的图像，并将其添加到结果列表中。
        x, y, w, h = box
        # 提取每个数字的图像
        digit_img = processed_image[y:y + h, x:x + w]
        digit_imgs.append(digit_img)
        if save_intermediate: #如果需要保存中间结果，将绘制了边界框的图像保存到指定目录。
            # 保存原始的数字图像
            cv2.imwrite(
                os.path.join(image_output_dir, f"{image_name}_2_{idx + 1}_segmentation.png"),
                digit_img,
            )

    # if save_intermediate:
    #     # 保存带有边界框的图像
    #     cv2.imwrite(
    #         os.path.join(image_output_dir, f"{image_name}_2_segmentation.png"), image_with_boxes
    #     )

    return digit_imgs #返回分割后的数字图像列表。


# 处理和归一化数字图像
def process_and_normalize_digits(
        digit_imgs, save_intermediate=False, output_dir=None, image_name=None
):
    image_name = os.path.splitext(os.path.basename(image_name))[0]
    if save_intermediate:
        # 创建保存中间结果的目录
        image_output_dir = os.path.join(output_dir, image_name)
        os.makedirs(image_output_dir, exist_ok=True)
    else:
        image_output_dir = None
    processed_digits = [] #初始化一个空列表，用于存储处理和归一化后的数字图像。
    for idx, digit in enumerate(digit_imgs): #遍历输入的数字图像列表。
        h, w = digit.shape #获取每个图像的高度和宽度
        if h == 0 or w == 0:
            # print(f"Warning: Skipping digit {idx + 1} because it's empty.")
            continue

        # 调整尺寸，保持纵横比，将图像缩放到20x20以内
        if h > w:
            new_h = 20
            new_w = int(w * (20 / h))
        else:
            new_w = 20
            new_h = int(h * (20 / w))

        # 防止新尺寸为0
        if new_w == 0:
            new_w = 1
        if new_h == 0:
            new_h = 1

        # 调整图像尺寸
        resized_digit = cv2.resize(digit, (new_w, new_h), interpolation=cv2.INTER_AREA)

        # 创建28x28的空白图像，并将调整后的数字放在中心位置
        #将调整大小后的数字图像放置到一个 28x28 的空白图像中心位置。
        padded_digit = np.zeros((28, 28), dtype=np.uint8)
        x_offset = (28 - new_w) // 2
        y_offset = (28 - new_h) // 2
        padded_digit[
        y_offset:y_offset + new_h, x_offset:x_offset + new_w
        ] = resized_digit

        # 计算笔画比例，用于判断是否需要调整笔画粗细
        non_zero_pixels = np.count_nonzero(padded_digit)
        total_pixels = padded_digit.size
        stroke_ratio = non_zero_pixels / total_pixels

        adjustment_type = "None"

        # 如果笔画过细，进行膨胀操作
        if stroke_ratio < 0.05:
            kernel_size = 2  # 根据需要调整内核大小
            kernel = np.ones((kernel_size, kernel_size), np.uint8)
            processed_digit = cv2.dilate(padded_digit, kernel)
            adjustment_type = f"Dilated with kernel size {kernel_size}"
        # 如果笔画过粗，进行腐蚀操作
        elif stroke_ratio > 0.3:
            kernel_size = 2  # 根据需要调整内核大小
            kernel = np.ones((kernel_size, kernel_size), np.uint8)
            processed_digit = cv2.erode(padded_digit, kernel)
            adjustment_type = f"Eroded with kernel size {kernel_size}"
        else:
            processed_digit = padded_digit

        if adjustment_type != "None":
            # 打印调整信息
            pass
            # print(f"Digit {idx + 1}: {adjustment_type} (stroke ratio: {stroke_ratio:.3f})")

        # 使用图像矩将数字中心化，通过计算图像的质心（重心），将数字图像的重心移动到图像中心。
        M = cv2.moments(processed_digit)
        if M["m00"] != 0:
            # 计算重心
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
            rows, cols = processed_digit.shape
            # 计算偏移量，使图像中心与重心对齐
            shiftx = (cols // 2) - cx
            shifty = (rows // 2) - cy
            M_shift = np.float32([[1, 0, shiftx], [0, 1, shifty]])
            # 应用仿射变换进行平移，使用最近邻插值
            centered_digit = cv2.warpAffine(
                processed_digit, M_shift, (cols, rows), flags=cv2.INTER_NEAREST
            )
        else:
            centered_digit = processed_digit

        # # **保存二值化前的图像**
        # if save_intermediate:
        #     cv2.imwrite(
        #         os.path.join(image_output_dir, f"{image_name}_digit_{idx + 1}_before_binarization.png"),
        #         centered_digit,
        #     )

        # 确保像素值为0或255，再次二值化（手动方式）
        threshold_value = 10  # 可以根据需要调整阈值
        _, centered_digit_binarized = cv2.threshold(centered_digit, threshold_value, 255, cv2.THRESH_BINARY)

        # **保存二值化后的图像**
        if save_intermediate:
            cv2.imwrite(
                os.path.join(image_output_dir, f"{image_name}_3_{idx + 1}_normalization.png"),
                centered_digit_binarized,
            )

        # 检查像素值是否仅为0或255
        unique_values = np.unique(centered_digit_binarized)
        # print(f"Digit {idx + 1} unique pixel values after binarization: {unique_values}")

        # 归一化像素值到[0, 1]，符合模型输入的要求。
        #归一化：将像素值从 [0, 255] 转换为 [0, 1]，便于模型处理。
        normalized_digit = centered_digit_binarized / 255.0
        processed_digits.append(normalized_digit)

    return processed_digits #返回处理和归一化后的数字图像列表。


# 识别数字
def recognize_digits(model, processed_digits, device):
    results = [] #创建一个空列表，用于存储每个数字的识别结果。
    model.eval()  # 确保模型在评估模式，确保模型在推理时的行为符合预期。
    with torch.no_grad(): #在推理阶段，不需要计算梯度（因为不会进行反向传播）
        for digit in processed_digits:
            # 将numpy数组转换为PyTorch张量，并添加批次和通道维度
            #将处理后的数字图像（NumPy 数组）转换为 PyTorch 张量，并调整其形状以符合模型输入要求。
            tensor_digit = (
                torch.from_numpy(digit).unsqueeze(0).unsqueeze(0).float().to(device)
            )
            # 标准化张量，与训练时的归一化方式保持一致
            tensor_digit = (tensor_digit - 0.5) / 0.5
            # 前向传播，获取模型输出，模型的输出是一个张量，表示每个类别的预测概率
            output = model(tensor_digit)
            # 获取预测的类别
            _, predicted = torch.max(output.data, 1)
            results.append(predicted.item())
    

    return results #返回一个包含所有数字识别结果的列表。


# 处理单张图像并返回识别结果
def process_image(
        image_path, model, device, save_intermediate=False, intermediate_dir=None
):
    try:#从输入图像路径中提取图像文件名（不包含扩展名）。recognize_digits
        image_name = os.path.splitext(os.path.basename(image_path))[0]

        if save_intermediate:
            # 创建保存中间结果的目录
            os.makedirs(intermediate_dir, exist_ok=True)
            image_output_dir = os.path.join(intermediate_dir, image_name)
            os.makedirs(image_output_dir, exist_ok=True)
        else:
            image_output_dir = None

        # 预处理图像，调用 preprocess_image 函数，对输入图像进行预处理。
        processed_image = preprocess_image(
            image_path, save_intermediate, image_output_dir, image_name
        )
        # 分割数字
        digit_imgs = segment_digits(
            processed_image, save_intermediate, image_output_dir, image_name
        )
        # 处理和归一化数字图像
        processed_digits = process_and_normalize_digits(
            digit_imgs, save_intermediate, image_output_dir, image_name
        )
        # 识别数字
        results = recognize_digits(model, processed_digits, device)
        # 将识别结果转换为字符串
        number_str = "".join(map(str, results))
        # print(f"识别结果 (在 process_image 中): {number_str}")
        return number_str

    except Exception as e:
        print(f"处理图像 {image_path} 时出错：{e}")
        return None
